import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt
import joblib  # 新增：用于模型保存和加载
# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


# 液态神经网络类
class LiquidNeuralNetwork:
    def __init__(self, n_nodes=200, input_dim=20, reservoir_density=0.5, spectral_radius=0.9):
        self.n_nodes = n_nodes  # 液态层神经元数量
        self.input_dim = input_dim  # 输入维度
        self.reservoir_density = reservoir_density  # 液态层连接密度
        self.spectral_radius = spectral_radius  # 液态层的谱半径
        
        # 初始化液态层的权重
        self.W_in = np.random.randn(self.n_nodes, self.input_dim) * 0.1  # 输入到液态层的权重
        self.W_reservoir = np.random.randn(self.n_nodes, self.n_nodes)  # 液态层内部连接权重
        # 调整谱半径
        max_eig = np.max(np.abs(np.linalg.eigvals(self.W_reservoir)))
        self.W_reservoir *= self.spectral_radius / max_eig
        
        # 液态层的节点状态
        self.states = np.zeros((self.n_nodes, 1))
        
        # 初始化读出层（简单起见，使用线性回归）
        self.readout = LinearRegression()
        
    def _update_states(self, input_data):
        # 更新液态层的状态
        input_effect = np.dot(self.W_in, input_data.reshape(-1, 1))
        reservoir_effect = np.dot(self.W_reservoir, self.states)
        self.states = np.tanh(input_effect + reservoir_effect)
        return self.states.flatten()
    
    def train(self, X, y):
        # 训练液态神经网络
        n_samples = X.shape[0]
        all_states = []
        
        for i in range(n_samples):
            # 更新液态层状态
            self._update_states(X[i])
            # 保存状态用于训练读出层
            all_states.append(self.states.flatten())
        
        # 将状态转换为 numpy 数组
        all_states = np.array(all_states)
        
        # 训练读出层
        self.readout.fit(all_states, y)
    
    def predict(self, X):
        # 预测未来值
        n_samples = X.shape[0]
        all_states = []
        
        for i in range(n_samples):
            # 更新液态层状态
            self._update_states(X[i])
            # 保存状态用于预测
            all_states.append(self.states.flatten())
        
        # 将状态转换为 numpy 数组
        all_states = np.array(all_states)
        
        # 使用读出层进行预测
        return self.readout.predict(all_states)

# 训练和预测函数
def train_lnn(X_train, y_train):
    lnn = LiquidNeuralNetwork(n_nodes=100, input_dim=X_train.shape[1])
    lnn.train(X_train, y_train)
    return lnn

def predict_lnn(lnn, X_test):
    return lnn.predict(X_test)

def generate_time_series_data(n_samples, time_step=0.1):
    """生成连续时间序列数据"""
    t = np.arange(0, n_samples * time_step, time_step)
    # 生成一个包含多个频率成分的信号
    signal = np.sin(2 * np.pi * 0.5 * t) + 0.5 * np.sin(2 * np.pi * 2 * t) + 0.2 * np.random.randn(len(t))
    return t, signal

def predict_continuous_time(lnn, initial_data, n_steps):
    """进行连续时间预测"""
    predictions = []
    current_state = initial_data[-1].reshape(1, -1)  # 使用最后一个已知数据点作为初始状态
    
    for _ in range(n_steps):
        # 预测下一个时间点
        next_pred = lnn.predict(current_state)
        predictions.append(next_pred[0])
        # 更新当前状态用于下一次预测
        current_state = next_pred.reshape(1, -1)
    
    return np.array(predictions)

# 保存模型函数
def save_lnn_model(model, filename='lnn_model.pkl'):
    joblib.dump(model, filename)
    print(f"模型已保存为 {filename}")

# 加载模型函数
def load_lnn_model(filename='lnn_model.pkl'):
    model = joblib.load(filename)
    print(f"模型已从 {filename} 加载")
    return model


# 生成时间序列数据
n_samples = 1000
time_step = 0.1
t, signal = generate_time_series_data(n_samples, time_step)

# 按百分比划分训练和测试数据
train_ratio = 0.9  # 90% 作为训练集
n_train = int(len(signal) * train_ratio)
n_test = len(signal) - n_train

# 准备训练数据
X_train = signal[:n_train].reshape(-1, 1)  # 前90%作为训练数据
y_train = signal[1:n_train+1].reshape(-1, 1)  # 目标值是下一个时间点的值

# 准备测试数据
X_test = signal[n_train:-1].reshape(-1, 1)  # 剩余10%作为测试数据
y_test = signal[n_train+1:].reshape(-1, 1)  # 对应的目标值

# 训练LNN
lnn_model = train_lnn(X_train, y_train)
save_lnn_model(lnn_model, 'lnn_model.pkl')  # 训练后自动保存模型

# 加载模型
lnn_model_loaded = load_lnn_model('lnn_model.pkl')

# 进行连续时间预测（使用加载的模型）
n_future_steps = 50
future_predictions = predict_continuous_time(lnn_model_loaded, X_test, n_future_steps)

# 评估预测结果
mse = mean_squared_error(y_test[:n_future_steps], future_predictions)
r2 = r2_score(y_test[:n_future_steps], future_predictions)
print(f"预测的均方误差: {mse}")
print(f"预测的R2值: {r2}")

# 可视化结果
future_t = t[-1] + np.arange(1, n_future_steps+1) * time_step
plt.figure(figsize=(12, 6))
plt.plot(t[-100:], signal[-100:], label='实际值', color='blue')
plt.plot(future_t, future_predictions, label='预测值', color='red', linestyle='--')
plt.xlabel('时间')
plt.ylabel('信号值')
plt.title('连续时间预测结果')
plt.legend()
plt.grid(True)
plt.show()

